"""
Pinecone结合MNIST的KNN搜索演示
安装依赖: pip install pinecone scikit-learn numpy matplotlib

使用说明:
1. 注册Pinecone账号获取API密钥
2. 替换下方API_KEY为你的密钥
3. 免费用户需使用us-east-1区域的serverless配置

常见问题:
- IProgress警告: 升级jupyter和ipywidgets
- 密钥错误: 检查API_KEY是否正确
"""
from __future__ import annotations
import time
from typing import List
import numpy as np
from sklearn.datasets import load_digits
from pinecone import Pinecone, ServerlessSpec

# 配置参数
INDEX_ID = "mnist-index"
VECTOR_DIM = 64  # 8x8图像展开维度
SIMILARITY_METRIC = "euclidean"
UPLOAD_BATCH = 100
API_KEY = "pcsk_4EcLgW_MpufGtAtj64LC3kASF1U834xnJP1M12k6ihSTyfkBteSHcGQpdTmjLVZzMkPRdb"


def prepare_index(pc: Pinecone, idx_name: str) -> None:
    """确保索引存在（存在则删除重建）"""
    try:
        existing_indexes = set(pc.list_indexes().names())
    except Exception:
        try:
            existing_indexes = {i["name"] for i in pc.list_indexes()}
        except Exception:
            existing_indexes = set()

    if idx_name in existing_indexes:
        pc.delete_index(idx_name)
        time.sleep(2)  # 等待删除完成

    # 创建新索引
    pc.create_index(
        name=idx_name,
        dimension=VECTOR_DIM,
        metric=SIMILARITY_METRIC,
        spec=ServerlessSpec(cloud="aws", region="us-east-1")
    )
    time.sleep(5)  # 等待索引就绪


def create_vectors(X: np.ndarray, y: np.ndarray) -> List[dict]:
    """将数据集转换为Pinecone可接收的格式"""
    vector_list = []
    for idx in range(len(X)):
        vector_list.append({
            "id": str(idx),
            "values": X[idx].astype(float).tolist(),
            "metadata": {"Label": int(y[idx])}
        })
    return vector_list


def upload_in_batches(index, vectors: List[dict], batch_size: int = UPLOAD_BATCH) -> None:
    """分批上传向量到Pinecone"""
    for start_idx in range(0, len(vectors), batch_size):
        end_idx = min(start_idx + batch_size, len(vectors))
        index.upsert(vectors=vectors[start_idx:end_idx])


def create_test_digit_3() -> np.ndarray:
    """创建自定义的8x8数字3图案（0-255范围）"""
    return np.array([
        [0, 255, 255, 255, 255, 255, 0, 0],
        [0, 0,   0,   0,   0, 255, 255, 0],
        [0, 0,   0,   0,   0, 255, 255, 0],
        [0, 0,   0,   0, 255, 255, 0, 0],
        [0, 0,   0,   0,   0, 255, 255, 0],
        [0, 0,   0,   0,   0, 255, 255, 0],
        [0, 255, 255, 255, 255, 255, 0, 0],
        [0, 0,   0,   0,   0, 0,   0, 0],
    ], dtype=np.float32)


def normalize_pixel_range(arr: np.ndarray) -> np.ndarray:
    """将0-255范围转换为0-16"""
    return (arr / 255.0 * 16.0).astype(np.float32)


def perform_knn_search(index, query: List[float]) -> None:
    """执行KNN查询并打印结果"""
    response = index.query(vector=query, top_k=11, include_metadata=True)
    matches = response.get("matches", []) if isinstance(response, dict) else getattr(response, "matches", [])

    labels = []
    result_parts = []
    for match in matches:
        vec_id = match.get("id") if isinstance(match, dict) else getattr(match, "id", "")
        distance = match.get("score") if isinstance(match, dict) else getattr(match, "score", 0)
        metadata = match.get("metadata") if isinstance(match, dict) else getattr(match, "metadata", {})
        label = int(metadata["Label"]) if metadata and "Label" in metadata else -1
        labels.append(label)
        result_parts.append(f"id:{vec_id},距离:{distance},标签:{label}")

    if result_parts:
        print(" ".join(result_parts))
    if labels:
        prediction = np.argmax(np.bincount(labels)).item()
        print(f"预测数字: {prediction}")


def run_demo() -> None:
    """主函数：初始化并运行演示"""
    try:
        pinecone_instance = Pinecone(api_key=API_KEY)
    except Exception as e:
        error_msg = str(e)
        if "Invalid API Key" in error_msg or "401" in error_msg:
            print("[错误] API密钥无效，请检查")
        else:
            print(f"[错误] 初始化失败: {error_msg}")
        return

    # 准备索引
    prepare_index(pinecone_instance, INDEX_ID)
    index = pinecone_instance.Index(INDEX_ID)

    # 加载并处理数据
    digits = load_digits(n_class=10)
    X_data = digits.data.astype(np.float32)
    y_labels = digits.target.astype(int)

    # 上传数据
    vectors = create_vectors(X_data, y_labels)
    upload_in_batches(index, vectors)

    # 创建测试查询
    test_image = create_test_digit_3()
    query_vector = normalize_pixel_range(test_image).flatten().tolist()

    # 执行查询
    perform_knn_search(index, query_vector)


if __name__ == "__main__":
    run_demo()